package streamAPI.operators;

import lombok.extern.slf4j.Slf4j;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple1;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.testng.annotations.Test;
import streamAPI.source.MyNoParallelSource;

@Slf4j
public class DemoWithPartition {

	@Test
	public void t1() throws Exception {
		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		DataStreamSource<String> source = env.socketTextStream("localhost", 4444);
		source.rebalance().map(new RichMapFunction<String, String>() {
			@Override
			public String map(String s) throws Exception {
				log.warn("{} -> index-{}", s, getRuntimeContext().getIndexOfThisSubtask());
				return s;
			}
		}).print();
		env.execute();
	}

	/**
	 * 自定义分区方式 泛型中是入参Key的类型
	 * 这样写的话就只支持 2 个分区
	 */
	class MyPartitionOnly2 implements Partitioner<Long> {
		@Override
		public int partition(Long in, int i) {
			if (in % 2 == 0) {
				return 0;
			}
			return 1;
		}
	}

	@Test(description = "自定义分区方式")
	public void t2() throws Exception {
		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		DataStreamSource<Long> source = env.addSource(new MyNoParallelSource());

		DataStream<Tuple1<Long>> s1 = source.map(new MapFunction<Long, Tuple1<Long>>() {
			@Override
			public Tuple1<Long> map(Long input) throws Exception {
				return new Tuple1<>(input);
			}
		}).setParallelism(1);
		
		s1.partitionCustom(new MyPartitionOnly2(), new KeySelector<Tuple1<Long>, Long>() {
			@Override
			public Long getKey(Tuple1<Long> input) throws Exception {
				return input.f0;
			}
		}).map(new MapFunction<Tuple1<Long>, String>() {
			@Override
			public String map(Tuple1<Long> in) throws Exception {
				return Thread.currentThread().getName() + "===>" + in.f0;
			}
		}).setParallelism(4).print().setParallelism(1);
		env.execute();
	}
}
